package cmd

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	pb "gitee.com/zfd81/dbrs/proto/dbrspb"
	pb3 "github.com/cheggaaa/pb/v3"
	"github.com/spf13/cobra"
)

const (
	OneK = 1024
	OneM = 1024 * OneK
	OneG = 1024 * OneM
)

func NewLoadCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "load <database table file>",
		Short: "Loads data file",
		Run:   LoadCommandFunc,
	}
	return cmd
}

func LoadCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 3 {
		ExitWithError(ExitBadArgs, fmt.Errorf("load cmd requires database name, table name and data file path as its argument"))
	}
	db := args[0]
	tbl := args[1]
	path := args[2]

	info, err := os.Stat(path)
	if err != nil || info.IsDir() {
		ExitWithError(ExitBadArgs, fmt.Errorf("open %s: No such file", path))
	}

	file, err := os.Open(path)
	if err != nil {
		ExitWithError(ExitBadArgs, fmt.Errorf("Read file %s failed: %v", path, err.Error()))
	}
	defer file.Close()

	token, err := getToken()
	if err != nil {
		ExitWithError(ExitError, err)
	}

	client, err := BrokerSystemClient()
	if err != nil {
		ExitWithError(ExitBadConnection, err)
	}

	tableInfo, err := client.TableInfo(context.Background(), &pb.TableRequest{
		Token:    token,
		Database: db,
		Table:    tbl,
	})
	if err != nil {
		ExitWithError(ExitBadFeature, err)
	}

	request := &pb.LoadRequest{
		Database: db,
		Table:    tbl,
		Lines:    []*pb.Line{},
	}

	fileSize := info.Size()
	count := 100
	cnt := 0
	bar := pb3.Simple.Start(count)
	go func() {
		for i := 0; i < count-1; i++ {
			bar.Increment()
			cnt++
			time.Sleep(time.Duration(fileSize/(10*OneM)) * 17 * time.Millisecond)
		}
	}()

	PrimaryKeyIndex := -1
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if strings.TrimSpace(scanner.Text()) != "" {
			fields := strings.Split(scanner.Text(), ",")
			for i, name := range fields {
				if name == tableInfo.PrimaryKey {
					PrimaryKeyIndex = i
				}
			}
			if PrimaryKeyIndex < 0 {
				ExitWithError(ExitBadArgs, fmt.Errorf("Primary key column %s not found", tableInfo.PrimaryKey))
			}
			request.Names = fields
			request.PrimaryKeyIndex = int64(PrimaryKeyIndex)
			break
		}
	}
	startTime := time.Now() //开始时间
	dataStream := make(chan string, 2000)
	go func() {
		for scanner.Scan() {
			dataStream <- scanner.Text()
		}
		close(dataStream)
	}()

	index := 0
	stream, err := client.Load(context.Background())
	if err != nil {
		ExitWithError(ExitBadFeature, err)
	}
	for data := range dataStream {
		if strings.TrimSpace(data) != "" {
			fields := strings.Split(data, ",")
			if index > 1000 {
				{
					err = stream.Send(request)
					if err != nil {
						if err == io.EOF {
							_, err = stream.CloseAndRecv()
						}
						ExitWithError(ExitBadFeature, err)
					}
					index = 0
					request.Lines = []*pb.Line{}
				}
			}
			request.Lines = append(request.Lines, &pb.Line{
				Fields: fields,
			})
			index++
		}
	}

	err = stream.Send(request)
	if err != nil {
		if err == io.EOF {
			_, err = stream.CloseAndRecv()
		}
		ExitWithError(ExitBadFeature, err)
	}

	resp, err := stream.CloseAndRecv()
	if err != nil {
		ExitWithError(ExitBadFeature, err)
	}

	bar.Add(count - cnt)
	bar.Finish()
	fmt.Printf("Record Count: %d\n", resp.Cnt)
	fmt.Printf("Elapsed time: %v\n", time.Since(startTime))
}
